#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np

def gen_golden_data_simple():
    dtype = np.float32
    num_features = 1093    # 可根据需求调整
    # 属性参数（默认值）
    momentum = 0.9
    epsilon = 1e-5

     # 生成输入数据（1D张量，符合ND格式要求）
    local_mean = np.random.uniform(-1.0, 1.0, size=(num_features,)).astype(dtype)
    local_var = np.random.uniform(0.1, 2.0, size=(num_features,)).astype(dtype)  # 方差为正数
    global_mean = np.random.uniform(-1.0, 1.0, size=(num_features,)).astype(dtype)
    global_var = np.random.uniform(0.1, 2.0, size=(num_features,)).astype(dtype)  # 方差为正数

    # 计算真值（golden data）：根据算子数学公式
    golden_updated_mean = momentum * global_mean + (1 - momentum) * local_mean
    golden_updated_var = momentum * global_var + (1 - momentum) * local_var + epsilon

    # 确保输入输出目录存在
    import os
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)

    # 写入二进制文件
    local_mean.tofile("./input/local_mean.bin")
    local_var.tofile("./input/local_var.bin")
    global_mean.tofile("./input/global_mean.bin")
    global_var.tofile("./input/global_var.bin")
    golden_updated_mean.tofile("./output/golden_updated_mean.bin")
    golden_updated_var.tofile("./output/golden_updated_var.bin")

    print(f"数据生成完成:")
    print(f"特征数量: {num_features}")
    print(f"输入文件: ./input/[local_mean|local_var|global_mean|global_var].bin")
    print(f"真值文件(golden): ./output/[golden_updated_mean|golden_updated_var].bin")


if __name__ == "__main__":
    gen_golden_data_simple()

